import os
import logging
import pandas as pd
import torch
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from transformers import BertTokenizer, BertForSequenceClassification
from tqdm.notebook import tqdm

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

# 配置参数
MAX_SEQ_LENGTH = 128
SENTIMENT_LABELS = ['负面', '中性', '正面']
DEFAULT_MODEL_PATH = './finbert_finetuned/finbert_finetuned_epoch_3'

class SentimentPredictor:
    """情感预测器类，用于预测文本的情感"""
    
    def __init__(self, model_path=DEFAULT_MODEL_PATH, max_seq_length=MAX_SEQ_LENGTH):
        """
        初始化情感预测器
        
        参数:
            model_path: 模型路径
            max_seq_length: 最大序列长度
        """
        self.model_path = model_path
        self.max_seq_length = max_seq_length
        
        # 检查CUDA是否可用
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        print(f"使用设备: {self.device}")
        
        # 加载分词器和模型
        self._load_tokenizer_and_model()
    
    def _load_tokenizer_and_model(self):
        """加载分词器和模型"""
        print(f"从 {self.model_path} 加载分词器和模型...")
        
        try:
            self.tokenizer = BertTokenizer.from_pretrained(self.model_path)
            self.model = BertForSequenceClassification.from_pretrained(self.model_path)
            self.model.to(self.device)
            print("分词器和模型加载成功")
        except Exception as e:
            print(f"模型加载失败: {str(e)}")
            raise
    
    def predict(self, texts, batch_size=16):
        """
        预测文本的情感
        
        参数:
            texts: 字符串或字符串列表，要预测的文本
            batch_size: 批处理大小
            
        返回:
            字典或字典列表，包含预测结果
        """
        # 如果输入是单个字符串，则转换为列表
        if isinstance(texts, str):
            texts = [texts]
            single_input = True
        else:
            single_input = False
        
        results = []
        
        # 批处理预测
        for i in range(0, len(texts), batch_size):
            batch_texts = texts[i:i+batch_size]
            batch_results = self._predict_batch(batch_texts)
            results.extend(batch_results)
        
        # 如果输入是单个字符串，则返回单个结果
        return results[0] if single_input else results
    
    def _predict_batch(self, texts):
        """
        批量预测文本的情感
        
        参数:
            texts: 字符串列表，要预测的文本
            
        返回:
            字典列表，包含预测结果
        """
        self.model.eval()
        
        # 对文本进行编码
        inputs = self.tokenizer(
            texts,
            padding=True,
            truncation=True,
            max_length=self.max_seq_length,
            return_tensors="pt"
        )
        
        # 将输入移动到设备上
        inputs = {k: v.to(self.device) for k, v in inputs.items()}
        
        # 预测
        with torch.no_grad():
            outputs = self.model(**inputs)
            logits = outputs.logits
            probabilities = torch.nn.functional.softmax(logits, dim=1)
        
        # 将结果转为numpy数组
        probs_numpy = probabilities.cpu().numpy()
        
        # 获取预测结果
        results = []
        for i, text in enumerate(texts):
            probs = probs_numpy[i]
            sentiment_idx = np.argmax(probs)
            sentiment = SENTIMENT_LABELS[sentiment_idx]
            confidence = probs[sentiment_idx]
            
            result = {
                'text': text,
                'predicted_sentiment': sentiment,
                'confidence': confidence,
                'negative_prob': probs[0],
                'neutral_prob': probs[1],
                'positive_prob': probs[2]
            }
            results.append(result)
        
        return results
    
    def predict_from_csv(self, input_file, text_column, output_file=None, batch_size=32):
        """
        从CSV文件中预测文本的情感
        
        参数:
            input_file: 输入CSV文件路径
            text_column: 文本列名
            output_file: 输出CSV文件路径，如果为None，则使用输入文件名+_predicted.csv
            batch_size: 批处理大小
            
        返回:
            DataFrame: 包含预测结果的DataFrame
        """
        print(f"从CSV文件 {input_file} 预测情感...")
        
        try:
            # 读取CSV文件
            df = pd.read_csv(input_file)
            print(f"CSV文件读取成功，共 {len(df)} 行")
            
            # 检查文本列是否存在
            if text_column not in df.columns:
                raise ValueError(f"列 '{text_column}' 不存在于CSV文件中")
            
            # 获取文本
            texts = df[text_column].tolist()
            
            # 预测情感
            print("开始预测情感...")
            results = self.predict(texts, batch_size)
            
            # 将结果添加到DataFrame
            df['predicted_sentiment'] = [r['predicted_sentiment'] for r in results]
            df['confidence'] = [r['confidence'] for r in results]
            df['negative_prob'] = [r['negative_prob'] for r in results]
            df['neutral_prob'] = [r['neutral_prob'] for r in results]
            df['positive_prob'] = [r['positive_prob'] for r in results]
            
            # 保存结果
            if output_file is None:
                base_name = os.path.splitext(input_file)[0]
                output_file = f"{base_name}_predicted.csv"
            
            df.to_csv(output_file, index=False)
            print(f"预测结果已保存到 {output_file}")
            
            # 可视化结果
            self.visualize_results(df)
            
            return df
        
        except Exception as e:
            print(f"从CSV文件预测失败: {str(e)}")
            raise
    
    def visualize_results(self, df):
        """
        可视化预测结果
        
        参数:
            df: 包含预测结果的DataFrame
        """
        print("可视化预测结果...")
        
        # 设置图表风格
        sns.set_style("whitegrid")
        
        # 绘制情感分布
        plt.figure(figsize=(12, 5))
        
        # 情感分布
        plt.subplot(1, 2, 1)
        ax = sns.countplot(x='predicted_sentiment', data=df, palette='viridis')
        plt.title('情感分布')
        plt.xlabel('情感类别')
        plt.ylabel('数量')
        
        # 添加百分比标签
        total = len(df)
        for p in ax.patches:
            height = p.get_height()
            ax.text(p.get_x() + p.get_width()/2., height + 0.1,
                    f'{height / total:.1%}',
                    ha="center")
        
        # 置信度分布
        plt.subplot(1, 2, 2)
        sns.histplot(df['confidence'], bins=20, kde=True, color='green')
        plt.title('置信度分布')
        plt.xlabel('置信度')
        plt.ylabel('频率')
        
        plt.tight_layout()
        plt.show()
        
        # 绘制置信度箱线图
        plt.figure(figsize=(8, 5))
        sns.boxplot(x='predicted_sentiment', y='confidence', data=df, palette='viridis')
        plt.title('不同情感类别的置信度分布')
        plt.xlabel('情感类别')
        plt.ylabel('置信度')
        plt.tight_layout()
        plt.show()


def load_predictor(model_path=DEFAULT_MODEL_PATH):
    """
    加载情感预测器
    
    参数:
        model_path: 模型路径
        
    返回:
        SentimentPredictor: 预测器实例
    """
    return SentimentPredictor(model_path=model_path)


def predict_text(text, predictor=None):
    """
    预测单个文本的情感
    
    参数:
        text: 要预测的文本
        predictor: 预测器实例，如果为None则创建新实例
        
    返回:
        dict: 预测结果
    """
    if predictor is None:
        predictor = load_predictor()
    
    result = predictor.predict(text)
    
    # 打印结果
    print(f"文本: {text}")
    print(f"预测情感: {result['predicted_sentiment']}")
    print(f"置信度: {result['confidence']:.4f}")
    print(f"情感概率分布:")
    print(f"  负面: {result['negative_prob']:.4f}")
    print(f"  中性: {result['neutral_prob']:.4f}")
    print(f"  正面: {result['positive_prob']:.4f}")
    
    return result


def predict_csv(input_file, text_column, output_file=None, model_path=DEFAULT_MODEL_PATH):
    """
    从CSV文件预测情感
    
    参数:
        input_file: 输入CSV文件路径
        text_column: 文本列名
        output_file: 输出CSV文件路径
        model_path: 模型路径
        
    返回:
        DataFrame: 包含预测结果的DataFrame
    """
    predictor = load_predictor(model_path)
    return predictor.predict_from_csv(input_file, text_column, output_file)


def batch_predict(texts, model_path=DEFAULT_MODEL_PATH):
    """
    批量预测文本情感
    
    参数:
        texts: 文本列表
        model_path: 模型路径
        
    返回:
        list: 预测结果列表
    """
    predictor = load_predictor(model_path)
    results = predictor.predict(texts)
    
    # 将结果转换为DataFrame并显示
    results_df = pd.DataFrame(results)
    display(results_df[['text', 'predicted_sentiment', 'confidence']])
    
    # 可视化结果
    predictor.visualize_results(results_df)
    
    return results 